#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
文档处理器模块
负责处理不同格式的文档并提取文本内容
"""

import os
import re
import logging
from typing import Dict, Any, List
from docx import Document
from PyPDF2 import PdfReader
import markdown
from bs4 import BeautifulSoup
import html2text

# 配置日志
logger = logging.getLogger(__name__)

class DocumentProcessor:
    """文档处理器类"""
    
    def __init__(self, config: Dict[str, Any] = None):
        """
        初始化文档处理器
        
        Args:
            config (dict): 文档处理器配置
        """
        self.config = config or {}
        
        # 默认配置
        self.default_config = {
            '支持格式': ['txt', 'pdf', 'docx', 'md', 'html'],
            '最大大小': 10485760  # 10MB
        }
        
        # 合并配置
        for key, value in self.default_config.items():
            if key not in self.config:
                self.config[key] = value
        
        # HTML转文本工具
        self.html_converter = html2text.HTML2Text()
        self.html_converter.ignore_links = True
        self.html_converter.ignore_images = True
        self.html_converter.ignore_emphasis = True
        
        logger.info(f"文档处理器初始化完成，支持格式：{self.config['支持格式']}")
    
    def update_config(self, config: Dict[str, Any]):
        """
        更新配置
        
        Args:
            config (dict): 新的配置
        """
        if config:
            self.config.update(config)
            logger.info(f"文档处理器配置已更新，支持格式：{self.config['支持格式']}")
    
    def extract_text(self, file_path: str, file_type: str = None) -> str:
        """
        从文档中提取文本
        
        Args:
            file_path (str): 文件路径
            file_type (str, optional): 文件类型，如果未提供则从文件扩展名推断
            
        Returns:
            str: 提取的文本内容
        """
        try:
            # 检查文件是否存在
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"文件不存在：{file_path}")
            
            # 检查文件大小
            file_size = os.path.getsize(file_path)
            if file_size > self.config['最大大小']:
                raise ValueError(f"文件大小超过限制：{file_size} > {self.config['最大大小']}")
            
            # 如果未提供文件类型，从文件扩展名推断
            if file_type is None:
                file_type = os.path.splitext(file_path)[1][1:].lower()
            
            # 检查文件类型是否支持
            if file_type not in self.config['支持格式']:
                raise ValueError(f"不支持的文件类型：{file_type}")
            
            # 根据文件类型选择相应的处理方法
            if file_type == 'txt':
                return self._process_txt(file_path)
            elif file_type == 'pdf':
                return self._process_pdf(file_path)
            elif file_type == 'docx':
                return self._process_docx(file_path)
            elif file_type == 'md':
                return self._process_markdown(file_path)
            elif file_type == 'html':
                return self._process_html(file_path)
            else:
                raise ValueError(f"未知的文件类型：{file_type}")
        
        except Exception as e:
            logger.error(f"文档处理失败: {str(e)}")
            raise
    
    def _process_txt(self, file_path: str) -> str:
        """
        处理TXT文件
        
        Args:
            file_path (str): 文件路径
            
        Returns:
            str: 提取的文本内容
        """
        try:
            # 尝试不同的编码方式
            encodings = ['utf-8', 'gbk', 'gb2312', 'ascii']
            
            for encoding in encodings:
                try:
                    with open(file_path, 'r', encoding=encoding) as f:
                        text = f.read()
                    return self._clean_text(text)
                except UnicodeDecodeError:
                    continue
            
            # 如果所有编码都失败，使用二进制模式读取
            with open(file_path, 'rb') as f:
                text = f.read().decode('utf-8', errors='ignore')
            return self._clean_text(text)
        
        except Exception as e:
            logger.error(f"TXT文件处理失败: {str(e)}")
            raise
    
    def _process_pdf(self, file_path: str) -> str:
        """
        处理PDF文件
        
        Args:
            file_path (str): 文件路径
            
        Returns:
            str: 提取的文本内容
        """
        try:
            # 打开PDF文件
            with open(file_path, 'rb') as f:
                # 创建PDF阅读器
                reader = PdfReader(f)
                
                # 提取所有页面的文本
                text = ''
                for page in reader.pages:
                    text += page.extract_text() + '\n'
            
            return self._clean_text(text)
        
        except Exception as e:
            logger.error(f"PDF文件处理失败: {str(e)}")
            raise
    
    def _process_docx(self, file_path: str) -> str:
        """
        处理DOCX文件
        
        Args:
            file_path (str): 文件路径
            
        Returns:
            str: 提取的文本内容
        """
        try:
            # 打开DOCX文件
            doc = Document(file_path)
            
            # 提取所有段落的文本
            text = '\n'.join([paragraph.text for paragraph in doc.paragraphs])
            
            return self._clean_text(text)
        
        except Exception as e:
            logger.error(f"DOCX文件处理失败: {str(e)}")
            raise
    
    def _process_markdown(self, file_path: str) -> str:
        """
        处理Markdown文件
        
        Args:
            file_path (str): 文件路径
            
        Returns:
            str: 提取的文本内容
        """
        try:
            # 读取Markdown文件
            with open(file_path, 'r', encoding='utf-8') as f:
                md_text = f.read()
            
            # 将Markdown转换为HTML
            html = markdown.markdown(md_text)
            
            # 将HTML转换为纯文本
            text = self.html_converter.handle(html)
            
            return self._clean_text(text)
        
        except Exception as e:
            logger.error(f"Markdown文件处理失败: {str(e)}")
            raise
    
    def _process_html(self, file_path: str) -> str:
        """
        处理HTML文件
        
        Args:
            file_path (str): 文件路径
            
        Returns:
            str: 提取的文本内容
        """
        try:
            # 读取HTML文件
            with open(file_path, 'r', encoding='utf-8') as f:
                html = f.read()
            
            # 使用BeautifulSoup解析HTML
            soup = BeautifulSoup(html, 'html.parser')
            
            # 删除script和style标签
            for script in soup(['script', 'style']):
                script.decompose()
            
            # 获取文本内容
            text = soup.get_text()
            
            return self._clean_text(text)
        
        except Exception as e:
            logger.error(f"HTML文件处理失败: {str(e)}")
            raise
    
    def _clean_text(self, text: str) -> str:
        """
        清理文本
        
        Args:
            text (str): 原始文本
            
        Returns:
            str: 清理后的文本
        """
        if not text:
            return ""
        
        # 替换多个换行符为单个换行符
        text = re.sub(r'\n+', '\n', text)
        
        # 替换多个空格为单个空格
        text = re.sub(r'\s+', ' ', text)
        
        # 删除开头和结尾的空白字符
        text = text.strip()
        
        return text
    
    def get_supported_formats(self) -> List[str]:
        """
        获取支持的文件格式
        
        Returns:
            list: 支持的文件格式列表
        """
        return self.config['支持格式']
    
    def get_max_file_size(self) -> int:
        """
        获取最大文件大小限制
        
        Returns:
            int: 最大文件大小（字节）
        """
        return self.config['最大大小']
    
    def format_file_size(self, size: int) -> str:
        """
        格式化文件大小
        
        Args:
            size (int): 文件大小（字节）
            
        Returns:
            str: 格式化后的文件大小
        """
        for unit in ['B', 'KB', 'MB', 'GB']:
            if size < 1024:
                return f"{size:.1f} {unit}"
            size /= 1024
        return f"{size:.1f} TB"